# coding:utf-8
# __user__ = hiicy redldw
# __time__ = 2019/8/8
# __file__ = data_gen
# __desc__ =
from pathlib import Path
import cv2
import pandas as pd
import numpy as np
from scipy import misc
import PIL


ftrain = r'f:\Resources\kdata\severstal-steel-defect-detection\train.csv'
save_masks = r'f:\Resources\kdata\severstal-steel-defect-detection\mask'
data = pd.read_csv(ftrain)
image_dir = r'f:\Resources\kdata\severstal-steel-defect-detection\train_images'
images = [img for img in Path(image_dir).iterdir()]
masks = [mask for mask in Path(save_masks).iterdir()]
def make_mask(row_id, df):
    '''Given a row index, return image_id and mask (256, 1600, 4)'''
    fname = df.iloc[row_id].name
    labels = df.iloc[row_id][:4]
    masks = np.zeros((256, 1600, 4), dtype=np.float32)  # float32 is V.Imp
    # 4:class 1～4 (ch:0～3)

    for idx, label in enumerate(labels.values):
        if label is not np.nan:
            label = label.split(" ")
            positions = map(int, label[0::2])
            length = map(int, label[1::2])
            mask = np.zeros(256 * 1600, dtype=np.uint8)
            for pos, le in zip(positions, length):
                mask[pos:(pos + le)] = 1
            masks[:, :, idx] = mask.reshape(256, 1600, order='F')
    return fname, masks


def gen_mask():
    for image in images:
        img = misc.imread(image)
        h,w = img.shape[:2]
        im_mask = np.zeros((h,w,4),dtype=np.int32) # todo: 用onehot变换
        image_name = Path(image).stem
        df = data.loc[data['ImageId_ClassId'].fillna("").str.contains(image_name)].reset_index()
        for i in range(len(df)):
            i_mask = np.zeros_like(im_mask[:,:,0].reshape(h*w),order='F')
            # one-hot
            pixels = df.loc[i,"EncodedPixels"]
            if type(pixels) is not str:
                continue
            pixells = pixels.split(" ")
            pixelems = [(int(pixells[pixel]), int(pixells[pixel + 1])) for pixel in range(0, len(pixells), 2)]
            pixeles = [pixeltuple[0] + i -1 for pixeltuple in pixelems for i in range(pixeltuple[1])]
            i_mask[pixeles] = 1
            im_mask[:, :, i] = i_mask.reshape((h,w),order="F")
        im_path = f'{save_masks}/{image_name}.png'
        cv2.imwrite(im_path,im_mask)

# txt_dir = r'f:\Resources\kdata\severstal-steel-defect-detection\index'
# txt_suffix = ['train','trainval','test']
# image_names = [Path(image).stem for image in images]
# limg = len(image_names)
# test_ratio = int(0.2*limg)
# trainval_ratio = int(0.2*limg)
# import random
# random.shuffle(image_names)
# test_txt = image_names[:test_ratio]
# trainval_txt = image_names[test_ratio:test_ratio+trainval_ratio]
# train_txt = image_names[test_ratio+trainval_ratio:]
#
# for suffix in txt_suffix:
#     with open(Path(txt_dir)/(suffix+".txt"),'w') as f:
#         for line in eval(suffix+"_txt"):
#             f.write(line+"\n")
# def func(x):
#     y=0
#     if type(x) is str:
#         y = 1
#     else:
#         y = 0
#     return y
# sss = {}
# data['ImageId_ClassId'] = data['ImageId_ClassId'].str.extract(".*\.jpg.(\d)")
# data['EncodedPixels'] = data['EncodedPixels'].apply(func)
# dj = data.groupby('ImageId_ClassId')['EncodedPixels'].sum()
# print(dj)

gen_mask()